import os
import json
from openai import OpenAI
from dotenv import load_dotenv
from serv.config import dblock
from dataclasses import asdict
from decimal import Decimal

load_dotenv()  # 加载环境变量

base_url = "https://dashscope.aliyuncs.com/compatible-mode/v1"
model = "qwen3-32b"
client = OpenAI(
    api_key=os.getenv("BAILIAN_KEY"),
    base_url=base_url,
)


def query_grade(stu_no: str) -> str:
    with dblock() as cur:
        cur.execute(
            """
            SELECT s.name as stu_name, c.name as course_name, g.grade 
            FROM course_grade as g
            INNER JOIN student as s ON g.stu_sn = s.sn
            INNER JOIN course as c  ON g.cou_sn = c.sn            
            WHERE s.no = %(stu_no)s 
            """,
            {"stu_no": stu_no},
        )

        items = [asdict(row) for row in cur]

        res = json.dumps(items, ensure_ascii=False, default=default_serializer)
        return res


# 定义工具（函数）描述
tools = [
    {
        "type": "function",
        "function": {
            "name": "query_grade",
            "description": "根据学号查询学生的成绩",
            "parameters": {
                "type": "object",
                "properties": {"stu_no": {"type": "str", "description": "学号"}},
                "required": ["stu_no"],
            },
        },
    },
]

system_prompt = r"""
你是一个用户的专属智能助手，负责处理用户请求。你的任务是：
1. 分析用户输入，确定请求与成绩相关。
2. 用户身份是李四（学号S002），你需要根据他的身份和请求，判断他是否有权限查询成绩。
3. 用户不能查询其它用户的成绩。
"""


def process_query(query: str) -> str:
    messages = [
        {"role": "assistant", "content": system_prompt},
        {"role": "user", "content": query},
    ]

    print(f"\n[1]REQ: {messages}\n")
    responses = client.chat.completions.create(
        model=model,
        messages=messages,
        tools=tools,
        tool_choice="auto",  # 让模型自动决定是否调用函数
        extra_body={"enable_thinking": False},
    )
    response = responses.choices[0]
    print(f"[1]RES: {response.model_dump()}\n")

    if response.finish_reason == "tool_calls":
        messages.append(response.message.model_dump())
        for tool_call in response.message.tool_calls:
            func_name = tool_call.function.name
            func_args = json.loads(tool_call.function.arguments)

            if func_name == "query_grade":
                result = query_grade(**func_args)
                messages.append(
                    {
                        "role": "tool",
                        "content": f"{result}",
                        "tool_call_id": tool_call.id,
                    }
                )

        print(f"[2]REQ: {messages}\n")
        responses = client.chat.completions.create(
            model=model,
            messages=messages,
            extra_body={"enable_thinking": False},
        )
        response = responses.choices[0]
        print(f"[2]RES: {response.model_dump()}\n")

        return response.message.content

    elif response.finish_reason == "stop":
        return response.message.content


def default_serializer(obj):
    if isinstance(obj, Decimal):
        return float(obj)  # 或者 str(obj)，取决于你希望如何表示
    raise TypeError(f"Object of type {type(obj)} is not JSON serializable")


if __name__ == "__main__":
    query = "我的课程是否都通过了？"
    print("问:", query)
    answer = process_query(query)
    print("答:", answer)
